from langchain.agents import AgentExecutor, create_tool_calling_agent, create_openai_tools_agent
from langchain.memory import ConversationBufferMemory
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_deepseek import ChatDeepSeek

import tool
import prompt_template
from simple_buffer_memory import SimpleConversationTokenBufferMemory


class Master(object):
    """
    黄大仙主类，完成算命、解梦等
    """
    def __init__(self):
        """
        初始化
        """

        self.mood = ""
        self.memory_key = "chat_history"

        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "{system_prompt}"),
            MessagesPlaceholder(variable_name=self.memory_key),
            ("human", "{question}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
            # MessagesPlaceholder(variable_name="tool_names"),
            # MessagesPlaceholder(variable_name="tools"),
        ])

        self.model = ChatDeepSeek(
            model="deepseek-chat",
            temperature=0,
            max_tokens=1024,
            # verbose=True
        )



        self.tools = [
            tool.empty_tools,
            tool.search_tools,
            tool.rag_tools,
            tool.bazi_cesuan,
        ]

        self.agent = create_tool_calling_agent(
            self.model,
            self.tools,
            self.prompt,
        )

        self.memory = self.get_memory()

        self.agent_executor = AgentExecutor(
            agent=self.agent,
            tools=self.tools,
            memory=self.memory,
            verbose=True,
            max_iterations=4,  # 防止无限循环，设置一个最大迭代次数
            early_stopping_method="force",  # 禁用需要 token 计数的 "generate" 方法
            # 关键：不要设置 handle_parsing_errors，或者用另一种方式处理
            # handle_parsing_errors=True 有时也会触发 token 检查
        )

    def get_memory(self):
        """
        获取记忆
        """
        chat_memory = RedisChatMessageHistory(
            session_id="session",
            url="redis://localhost:6379",

        )
        # chat_memory.clear()

        memory = SimpleConversationTokenBufferMemory(
            memory_key=self.memory_key,
            llm=self.model,
            human_prefix="用户",
            ai_prefix="黄大仙",
            return_messages=True,
            input_key="question",
            chat_memory=chat_memory
        )
        return memory

    def mood_judge(self, query):
        """
        情绪链，从用户输入判断用户情绪
        """
        chain = ChatPromptTemplate.from_template(prompt_template.mood_prompt_template) | self.model | StrOutputParser()
        result = chain.invoke({"mood_input": query})
        return result

    def run(self, query):
        """
        运行
        """
        self.mood = self.mood_judge(query)
        # print(self.mood)
        _role_set = prompt_template.mood_role_set[self.mood]["roleSet"]
        # print(_role_set)
        _main_prompt = prompt_template.main_prompt_template.format(roleSet=_role_set)
        # print(_main_prompt)
        return self.agent_executor.invoke({
            "question": query,
            "system_prompt": _main_prompt
        })
